import os

import h5py
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import rasterio
import rasterio.plot
from rasterio.io import MemoryFile
from rasterio.mask import mask
from rasterio.transform import from_origin
import geopandas as gpd
from shapely.geometry import mapping

from Timer import Timer


def bilinear_interpolation(x, y, values, invalid_value=-999):
    """
    对 (x, y) 使用四邻点的值进行双线性插值，同时考虑无效值。

    :param x: float, 目标点的x坐标
    :param y: float, 目标点的y坐标
    :param values: list, 四邻点的值，格式为[(x1, y1, f11), (x1, y2, f12), (x2, y1, f21), (x2, y2, f22)]
    :param invalid_value: float, 用于指示无效值的特殊值，默认为 -999
    :return: float, 插值结果
    """
    (x1, y1, f11), (x1, y2, f12), (x2, y1, f21), (x2, y2, f22) = values

    # 筛选出有效的邻点值
    valid_values = [(f, (x1, y1) if f == f11 else
    (x1, y2) if f == f12 else
    (x2, y1) if f == f21 else
    (x2, y2))
                    for f in (f11, f12, f21, f22) if f != invalid_value]

    # 如果所有邻点值都是无效的，返回无效值
    if not valid_values:
        return invalid_value

    # 初始化权重和有效值的累加器
    total_weight = 0
    weighted_sum = 0

    for f, (xi, yi) in valid_values:
        weight = (1 - abs(x - xi) / (x2 - x1)) * (1 - abs(y - yi) / (y2 - y1))
        weighted_sum += f * weight
        total_weight += weight

    return weighted_sum / total_weight if total_weight != 0 else invalid_value


class SatelliteH5Dataset:

    def __init__(self, dataset):
        self.name = dataset.name
        self.shape = dataset.shape
        self.width = dataset.shape[1]
        self.height = dataset.shape[0]
        self.ndim = dataset.ndim
        self.dtype = dataset.dtype
        self.fillValue = dataset.attrs.get("FillValue", -999)
        self.intercept = dataset.attrs.get("Intercept", 0)
        self.slope = dataset.attrs.get("Slope", 1)
        self.units = dataset.attrs.get("Units", '')
        if self.units == '':
            self.units = dataset.attrs.get("units", '')
        self.dataset = dataset
        self.data = np.array(dataset)
        self.valid_pixels_result = -1

    def valid_pixels(self):
        if self.valid_pixels_result < 0:
            self.valid_pixels_result = np.sum(self.data != self.fillValue)
        return self.valid_pixels_result

    def valid_rate(self):
        return self.valid_pixels() / self.data.size

    def pixel_to_latlon(self, x, y):
        if x < 0 or x >= self.width or y < 0 or y >= self.height:
            raise ValueError("输入参数超限")

        longitude = x / self.width * 360 - 180
        latitude = 90 - y / self.height * 180
        return latitude, longitude

    def latlon_to_pixel(self, lat, lon):
        if lat < -90 or lat >= 90 or lon < -180 or lon >= 180:
            raise ValueError("输入参数超限")

        x = (lon + 180) / 360 * self.width
        y = (90 - lat) / 180 * self.height
        return int(x), int(y)

    def get_value(self, lat, lon):
        return self.data[self.latlon_to_pixel(lat, lon)]

    def get_interpolated_value(self, lat, lon):
        if lat < -90 or lat >= 90 or lon < -180 or lon >= 180:
            raise ValueError("输入参数超限")

        x = (lon + 180) / 360 * self.width
        y = (90 - lat) / 180 * self.height

        # 获取四邻点的值
        x1, y1 = int(np.floor(x)), int(np.floor(y))
        x2, y2 = x1 + 1, y1 + 1
        # 保证不超出边界
        if x2 >= self.width:
            x2 = self.width - 1
        if y2 >= self.height:
            y2 = self.height - 1
        f11 = self.data[y1, x1]
        f12 = self.data[y2, x1]
        f21 = self.data[y1, x2]
        f22 = self.data[y2, x2]
        values = [(x1, y1, f11), (x1, y2, f12), (x2, y1, f21), (x2, y2, f22)]

        return bilinear_interpolation(x, y, values, self.fillValue)

    def par(self, par_func=None):
        if par_func is None:
            def par_func(ds):
                return ds * self.slope + self.intercept

        self.data = par_func(self.data)
        return self

    def dim_reduce(self, reduce_func=np.mean):
        if len(self.data.shape) <= 1:
            raise ValueError("无法降维")

        self.ndim -= 1
        self.data = reduce_func(self.data, axis=self.ndim)
        return self

    def subplot(self, ax, use_latlon_croods=True, interpolation=False):
        while self.ndim > 2:
            self.dim_reduce()
        ds = np.array(self.data)
        if ds.dtype.kind in 'iu':  # 如果数据类型是整数或无符号整数
            ds = ds.astype(float)
        ds[ds == self.fillValue] = np.nan

        if use_latlon_croods:
            lon = np.linspace(-180, 180, self.width)
            lat = np.linspace(90, -90, self.height)
            extent = [lon.min(), lon.max(), lat.min(), lat.max()]
            ax.set_xlabel("Longitude")
            ax.set_ylabel("Latitude")
        else:
            extent = None

        if not interpolation:
            interpolation = 'none'

        return ax.imshow(ds, cmap='viridis', aspect='auto', extent=extent, interpolation=interpolation)

    def plot(self, width=1000, height=500, use_latlon_croods=True, interpolation=False):
        px = 1 / plt.rcParams['figure.dpi']
        fig = plt.figure(figsize=(width * px, height * px))
        ax = fig.add_subplot(1, 1, 1)
        ax.set_title(self.name)
        im = self.subplot(ax, use_latlon_croods, interpolation)
        cbar = fig.colorbar(im, ax=ax)
        if len(self.units) > 0:
            cbar.set_label(self.units)
        plt.tight_layout()
        plt.show()


class SatelliteH5:

    def __init__(self, path):
        self.path = path
        self.file = h5py.File(path, 'r')

    def __del__(self):
        self.file.close()
        self.file = None

    def summary(self):
        def print_group(name, obj):
            # 计算缩进级别
            level = name.count('/')
            # 打印缩进和名称
            print('  ' * level + name)
            for key, val in obj.attrs.items():
                print('  ' * (level + 1) + '{}: {}'.format(key, val))

        self.file.visititems(print_group)

    def __getitem__(self, key):
        if isinstance(key, int) and key >= 0:
            key = list(self.file.keys())[key]
        group = self.file[key]
        if isinstance(group, h5py.Dataset):
            return SatelliteH5Dataset(group)
        return group

    def plot(self, width=1000, height=500, use_latlon_croods=True, interpolation=False):
        datasets = {}

        def check_dataset_2d(name, obj):
            if isinstance(obj, h5py.Dataset) and obj.ndim >= 2:
                datasets[name] = obj

        self.file.visititems(check_dataset_2d)
        n = len(datasets)
        px = 1 / plt.rcParams['figure.dpi']
        fig = plt.figure(figsize=(width * px, height * n * px))
        i = 1
        for key in datasets:
            h5ds = SatelliteH5Dataset(datasets[key])
            ax = fig.add_subplot(n, 1, i)
            ax.set_title(key)
            im = h5ds.subplot(ax, use_latlon_croods, interpolation)
            cbar = fig.colorbar(im, ax=ax)
            if len(h5ds.units) > 0:
                cbar.set_label(h5ds.units)
            i += 1
        plt.tight_layout()
        plt.show()

    def calc_and_save(self, map_to_calc, calc_func, filename):
        result = {}
        for key in map_to_calc:
            h5_ds0, h5_ds1 = map_to_calc[key]
            ds = calc_func(h5_ds0, h5_ds1)
            result[key] = (h5_ds0.shape, h5_ds0.dtype, ds)

        dirname = os.path.dirname(self.path)
        path = os.path.join(dirname, filename)
        f = h5py.File(path, "w")

        for key in result:
            shape, dtype, ds = result[key]
            f.create_dataset(key, shape, dtype, ds)

        f.flush()
        f.close()

        return SatelliteH5(path)

    def calc_file(self, other, calc_func, filename):
        if self.file.keys() != other.file.keys():
            raise ValueError("运算参数不匹配")

        map_to_calc = {}

        def check_dataset(name, obj):
            if isinstance(obj, h5py.Dataset) and isinstance(other.file[name], h5py.Dataset):
                h5_ds0 = SatelliteH5Dataset(obj)
                h5_ds1 = SatelliteH5Dataset(other.file[name])
                if h5_ds0.shape != h5_ds1.shape or h5_ds0.dtype != h5_ds1.dtype:
                    raise ValueError("运算参数不匹配")
                map_to_calc[name] = (h5_ds0, h5_ds1)

        self.file.visititems(check_dataset)

        return self.calc_and_save(map_to_calc, calc_func, filename)

    def __sub__(self, other):
        def sub_func(h5_ds0, h5_ds1):
            ds0 = np.array(h5_ds0.dataset)
            ds1 = np.array(h5_ds1.dataset)
            ds = np.empty_like(ds0)
            ds.fill(h5_ds0.fillValue)
            mask0 = ds0 != h5_ds0.fillValue
            mask1 = ds1 != h5_ds1.fillValue
            valid_mask = mask0 & mask1
            ds[valid_mask] = ds0[valid_mask] - ds1[valid_mask]
            return ds

        return self.calc_file(other, sub_func, "result.hdf5")

    def __or__(self, other):
        def merge(h5_ds0, h5_ds1):
            ds0 = np.array(h5_ds0.dataset)
            ds1 = np.array(h5_ds1.dataset)
            ds = np.empty_like(ds0)
            ds.fill(h5_ds0.fillValue)
            mask0 = ds0 != h5_ds0.fillValue
            ds[mask0] = ds0[mask0]
            mask1 = ds1 != h5_ds1.fillValue
            ds[mask1] = ds1[mask1]
            valid_mask = mask0 & mask1
            ds[valid_mask] = (ds0[valid_mask] + ds1[valid_mask]) / 2
            return ds

        return self.calc_file(other, merge, "result.hdf5")


def L1_globalize_EARTH_OBSERVE_BT(l1_files, output_file, calc_func=None, global_width=5552, global_height=2336):
    timer = Timer()
    timer.tick("L1_globalize_EARTH_OBSERVE_BT begin")
    ds_sample = SatelliteH5(l1_files[0])["Sounding Channel/Calibration/EARTH_OBSERVE_BT"]
    print(ds_sample.shape)
    ds_result = np.full((global_height, global_width), ds_sample.fillValue, ds_sample.dtype)

    if calc_func is None:
        def calc_func(ds0, ds1):
            valid_mask = ds0 == ds_sample.fillValue
            return np.where(valid_mask, ds1, (ds0 + ds1) / 2)

    def map_to_global(value, value_min, value_max, global_min, global_max):
        return np.clip(np.round((value - value_min) / (value_max - value_min) * (global_max - global_min)).astype(int),
                       global_min, global_max - 1)

    for l1_file in l1_files:
        filename = os.path.basename(l1_file)
        timer.tick(f"{filename} 开始处理")
        h5_l1 = SatelliteH5(l1_file)
        h5_l1_eob = h5_l1["Sounding Channel/Calibration/EARTH_OBSERVE_BT"]
        ds_value = h5_l1_eob.par().dim_reduce(np.mean).data.ravel()
        ds_lat = h5_l1["Sounding Channel/Geolocation/Latitude"].data.ravel()
        ds_lon = h5_l1["Sounding Channel/Geolocation/Longitude"].data.ravel()
        timer.tick(f"{filename} 读取 ds_value, ds_lat, ds_lon 为一维等长数组")

        global_x = map_to_global(ds_lon, -180, 180, 0, global_width)
        global_y = map_to_global(ds_lat, 90, -90, 0, global_height)
        timer.tick(f"{filename} 处理经纬度坐标转换")

        mask = ds_value != ds_sample.fillValue
        ds_result[global_y[mask], global_x[mask]] = calc_func(ds_result[global_y[mask], global_x[mask]], ds_value[mask])
        timer.tick(f"{filename} 处理有效值的叠加")
        timer.tick(f"{filename} 处理完成")

    f = h5py.File(output_file, "w")
    f_ds = f.create_dataset("EARTH_OBSERVE_BT_globalized", (global_height, global_width), ds_sample.dtype, ds_result)
    f_ds.attrs["FillValue"] = ds_sample.fillValue
    f_ds.attrs["Units"] = ds_sample.units
    f.flush()
    f.close()

    timer.tick("L1_globalize_EARTH_OBSERVE_BT done")
    return SatelliteH5(output_file)


def L2_merge_VSM(h5, filename="merged.hdf5"):
    map_to_calc = {
        "VSM_AD": (SatelliteH5Dataset(h5.file["VSM_A"]),
                   SatelliteH5Dataset(h5.file["VSM_D"]))
    }

    def merge(h5_ds0, h5_ds1):
        ds0 = np.array(h5_ds0.dataset)
        ds1 = np.array(h5_ds1.dataset)
        ds = np.empty_like(ds0)
        ds.fill(h5_ds0.fillValue)
        mask0 = ds0 != h5_ds0.fillValue
        ds[mask0] = ds0[mask0]
        mask1 = ds1 != h5_ds1.fillValue
        ds[mask1] = ds1[mask1]
        valid_mask = mask0 & mask1
        ds[valid_mask] = (ds0[valid_mask] + ds1[valid_mask]) / 2
        return ds

    return h5.calc_and_save(map_to_calc, merge, filename)


def L2_stat(directory, output_filename="stat.xlsx"):
    timer = Timer()
    timer.tick("L2_stat begin")
    columns = ['文件名', '宽', '高', 'VSM_A 有效点数', 'VSM_A 有效点占比', 'VSM_D 有效点数', 'VSM_A 有效点占比']
    results = []

    for filename in os.listdir(directory):
        if filename.lower().endswith('.hdf'):
            filepath = os.path.join(directory, filename)
            h5 = SatelliteH5(filepath)
            h5_vsm_a = h5["VSM_A"]
            h5_vsm_d = h5["VSM_D"]
            # 添加文件名和结果到列表
            results.append((
                filename,
                h5_vsm_a.width,
                h5_vsm_a.height,
                h5_vsm_a.valid_pixels(),
                h5_vsm_a.valid_rate(),
                h5_vsm_d.valid_pixels(),
                h5_vsm_d.valid_rate()
            ))
            timer.tick(filename, "stat complete")

    df = pd.DataFrame(results, columns=columns)
    print(df.head())

    output_xls = f"{directory}/{output_filename}"
    df.to_excel(output_xls, index=False)
    timer.tick("L2_stat done")


def L2_calc_all(input_xls, find_fy_files_func, output_xls):
    timer = Timer()
    timer.tick("L2_calc_all begin")
    columns = ['观测时间', '经度', '纬度', 'ASM', 'FY-3F_A', 'FY-3F_D', 'FY-3D_A', 'FY-3D_D']
    # 初始化一个空的DataFrame
    df = pd.DataFrame(columns=columns)
    df0 = pd.read_excel(input_xls)
    df0['日期'] = pd.to_datetime(df0['日期']).dt.strftime('%Y%m%d')
    df0 = df0.groupby(['站号', '日期', '经度', '纬度'])['土壤体积含水量'].mean().reset_index()

    df['观测时间'] = df0['日期']
    df['经度'] = df0['经度']
    df['纬度'] = df0['纬度']
    df['ASM'] = df0['土壤体积含水量']
    timer.tick("日期、经度、纬度、土壤体积含水量 就绪")

    timer.tick(f"begin filling interpolated values for {len(df)} rows")
    i = 0
    for index, row in df.iterrows():
        date = row['观测时间']
        longitude = row['经度']
        latitude = row['纬度']
        fy_3d, fy_3f = find_fy_files_func(date)
        df.at[index, 'FY-3F_A'] = fy_3f['VSM_A'].get_interpolated_value(latitude, longitude)
        df.at[index, 'FY-3F_D'] = fy_3f['VSM_D'].get_interpolated_value(latitude, longitude)
        df.at[index, 'FY-3D_A'] = fy_3d['VSM_A'].get_interpolated_value(latitude, longitude)
        df.at[index, 'FY-3D_D'] = fy_3d['VSM_D'].get_interpolated_value(latitude, longitude)
        i += 1
        if timer.elapsed(15):
            timer.tick(f"{i} rows filled by interpolated values")

    print(df.head())

    df.to_excel(output_xls, index=False)
    timer.tick("L2_calc_all done")


def L2_export_henan_tiff(h5ds, shapefile_path, output_file):
    ds = h5ds.data
    # 根据栅格数据的大小计算像素宽度和高度
    pixel_width = 360 / h5ds.width
    pixel_height = 180 / h5ds.height
    # 创建仿射变换
    transform = from_origin(-180, 90, pixel_width, pixel_height)
    # 创建元数据字典
    meta = {
        "driver": "GTiff",
        "count": 1,
        "dtype": "float64",
        "nodata": h5ds.fillValue,
        "width": h5ds.width,
        "height": h5ds.height,
        "crs": "+proj=latlong"
    }

    # 读取shapefile文件
    shapefile = gpd.read_file(shapefile_path)
    # 提取shapefile中的几何信息
    geometries = [mapping(geom) for geom in shapefile.geometry]

    # 使用 MemoryFile 创建临时 rasterio 数据集
    with MemoryFile() as memfile:
        with memfile.open(**meta, transform=transform) as dataset:
            dataset.write(ds, 1)

            # 使用蒙版裁剪栅格数据
            out_image, out_transform = mask(dataset, geometries, crop=True)
            out_meta = dataset.meta.copy()

            # 更新元数据，以适应裁剪后的数据
            out_meta.update({
                "driver": "GTiff",
                "nodata": h5ds.fillValue,
                "height": out_image.shape[1],
                "width": out_image.shape[2],
                "transform": out_transform
            })

            # 将裁剪后的栅格数据写入新的GeoTIFF文件
            with rasterio.open(output_file, "w", **out_meta) as dest:
                dest.write(out_image)


# 使用 rasterio 读取生成的 TIFF 文件并显示
def plot_tiff(tiff_path, invalid_value=-999):
    with rasterio.open(tiff_path) as src:
        fig, ax = plt.subplots(figsize=(10, 10))
        raster_data = src.read(1)

        # 将蒙版之外的值设置为 NaN
        raster_data[raster_data == invalid_value] = np.nan

        cax = ax.imshow(raster_data, cmap='viridis', extent=(src.bounds.left, src.bounds.right, src.bounds.bottom, src.bounds.top), interpolation='none')
        cbar = fig.colorbar(cax, ax=ax, orientation='vertical')
        cbar.set_label('Value')
        ax.set_title("Clipped Raster Data")
        ax.set_xlabel("Longitude")
        ax.set_ylabel("Latitude")
        plt.show()
